package com.zhongen.common.datascope.interceptor;

import cn.hutool.core.convert.Convert;
import cn.zhxu.bs.SearchSql;
import cn.zhxu.bs.SqlInterceptor;
import cn.zhxu.bs.param.FetchType;
import com.zhongen.common.datascope.DataScope;
import com.zhongen.common.datascope.handler.DataPermissionHandler;
import com.zhongen.common.datascope.holder.DataScopeMatchNumHolder;
import com.zhongen.common.datascope.holder.MappedStatementIdsWithoutDataScope;
import com.zhongen.common.datascope.processor.DataScopeSqlProcessor;
import com.zhongen.common.datascope.util.PluginUtils;
import lombok.RequiredArgsConstructor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;

import java.sql.Connection;
import java.util.List;
import java.util.Map;

/**
 * 数据权限拦截器
 */
@RequiredArgsConstructor
@Intercepts({
		@Signature(type = StatementHandler.class, method = "prepare", args = { Connection.class, Integer.class }) })
public class DataPermissionInterceptor implements Interceptor, SqlInterceptor {

	private final DataScopeSqlProcessor dataScopeSqlProcessor;

	private final DataPermissionHandler dataPermissionHandler;

	@Override
	public Object intercept(Invocation invocation) throws Throwable {
		// 第一版，测试用
		Object target = invocation.getTarget();
		StatementHandler sh = (StatementHandler) target;
		PluginUtils.MPStatementHandler mpSh = PluginUtils.mpStatementHandler(sh);
		MappedStatement ms = mpSh.mappedStatement();
		SqlCommandType sct = ms.getSqlCommandType();
		PluginUtils.MPBoundSql mpBs = mpSh.mPBoundSql();
		String mappedStatementId = ms.getId();

		// 获取当前需要控制的 dataScope 集合
		List<DataScope> filterDataScopes = dataPermissionHandler.filterDataScopes(mappedStatementId);
		if (filterDataScopes == null || filterDataScopes.isEmpty()) {
			return invocation.proceed();
		}

		// 根据用户权限判断是否需要拦截，例如管理员可以查看所有，则直接放行
		if (dataPermissionHandler.ignorePermissionControl(filterDataScopes, mappedStatementId)) {
			return invocation.proceed();
		}

		// 创建 matchNumTreadLocal
		DataScopeMatchNumHolder.initMatchNum();
		try {
			// 根据 DataScopes 进行数据权限的 sql 处理
			if (sct == SqlCommandType.SELECT) {
				mpBs.sql(dataScopeSqlProcessor.parserSingle(mpBs.sql(), filterDataScopes));
			}
			else if (sct == SqlCommandType.INSERT || sct == SqlCommandType.UPDATE || sct == SqlCommandType.DELETE) {
				mpBs.sql(dataScopeSqlProcessor.parserMulti(mpBs.sql(), filterDataScopes));
			}
			// 如果解析后发现当前 mappedStatementId 对应的 sql，没有任何数据权限匹配，则记录下来，后续可以直接跳过不解析
			Integer matchNum = DataScopeMatchNumHolder.pollMatchNum();
			List<DataScope> allDataScopes = dataPermissionHandler.dataScopes();
			if (allDataScopes.size() == filterDataScopes.size() && matchNum != null && matchNum == 0) {
				MappedStatementIdsWithoutDataScope.addToWithoutSet(filterDataScopes, mappedStatementId);
			}
		}
		finally {
			DataScopeMatchNumHolder.removeIfEmpty();
		}

		// 执行 sql
		return invocation.proceed();
	}

	@Override
	public <T> SearchSql<T> intercept(SearchSql<T> searchSql, Map<String, Object> map, FetchType fetchType) {
		String hashCode;
		final boolean shouldQueryCluster = fetchType.shouldQueryCluster();
		final boolean shouldQueryList = fetchType.shouldQueryList();
		String sqlString = shouldQueryCluster ? searchSql.getClusterSqlString() : searchSql.getListSqlString();
		hashCode = Convert.toStr(sqlString.hashCode());

		// 获取当前需要控制的 dataScope 集合
		List<DataScope> filterDataScopes = dataPermissionHandler.filterDataScopes(hashCode);
		if (filterDataScopes == null || filterDataScopes.isEmpty()) {
			return searchSql;
		}

		// 根据用户权限判断是否需要拦截，例如管理员可以查看所有，则直接放行
		if (dataPermissionHandler.ignorePermissionControl(filterDataScopes, hashCode)) {
			return searchSql;
		}

		if (shouldQueryCluster) {
			// 创建 matchNumTreadLocal
			DataScopeMatchNumHolder.initMatchNum();
			try {
				final String countSql = searchSql.getClusterSqlString();
				searchSql.setClusterSqlString(dataScopeSqlProcessor.parserSingle(countSql, filterDataScopes));

				// 根据 DataScopes 进行数据权限的 sql 处理
				// 如果解析后发现当前 hashCode 对应的 sql，没有任何数据权限匹配，则记录下来，后续可以直接跳过不解析
				Integer matchNum = DataScopeMatchNumHolder.pollMatchNum();
				List<DataScope> allDataScopes = dataPermissionHandler.dataScopes();
				if (allDataScopes.size() == filterDataScopes.size() && matchNum != null && matchNum == 0) {
					MappedStatementIdsWithoutDataScope.addToWithoutSet(filterDataScopes, hashCode);
				}
			}
			finally {
				DataScopeMatchNumHolder.removeIfEmpty();
			}

		}
		if (shouldQueryList) {
			// 创建 matchNumTreadLocal
			DataScopeMatchNumHolder.initMatchNum();
			try {
				final String listSql = searchSql.getListSqlString();
				searchSql.setListSqlString(dataScopeSqlProcessor.parserSingle(listSql, filterDataScopes));

				// 根据 DataScopes 进行数据权限的 sql 处理
				// 如果解析后发现当前 hashCode 对应的 sql，没有任何数据权限匹配，则记录下来，后续可以直接跳过不解析
				Integer matchNum = DataScopeMatchNumHolder.pollMatchNum();
				List<DataScope> allDataScopes = dataPermissionHandler.dataScopes();
				if (allDataScopes.size() == filterDataScopes.size() && matchNum != null && matchNum == 0) {
					MappedStatementIdsWithoutDataScope.addToWithoutSet(filterDataScopes, hashCode);
				}
			}
			finally {
				DataScopeMatchNumHolder.removeIfEmpty();
			}
		}
		return searchSql;
	}

	@Override
	public Object plugin(Object target) {
		if (target instanceof StatementHandler) {
			return Plugin.wrap(target, this);
		}
		return target;
	}

}
